/*
 * Copyright (c) 2023 Shenzhen Kaihong Digital Industry Development Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "rtmp_codec_h264.h"
#include <chrono>
#include <netinet/in.h>
#include "common/sharing_log.h"

namespace OHOS {
namespace Sharing {

bool RtmpEncoderH264::InputFrame(const uint8_t *frame, size_t size)
{
    if (!frame || !size) {
        return false;
    }

    const uint8_t *data = frame;
    size_t len = size;

    if (len > 4 && data[0] == 0 && data[1] == 0) {
        if (data[2] == 0x01) {
            data += 3;
            len -= 3;
        } else if (data[2] == 0 && data[3] == 0x01) {
            data += 4;
            len -= 4;
        }
    }

    uint8_t naluType = data[0] & 0x1f;
    if (naluType == NAL_SPS) {
        SHARING_LOGD("get SPS\n");
        sps_.assign((char *)data, (char *)data + len);
        return true;
    } else if (naluType == NAL_PPS) {
        SHARING_LOGD("get PPS\n");
        pps_.assign((char *)data, (char *)data + len);
        MakeConfigPacket();
        return true;
    } else if (naluType == NAL_SEI) {
        SHARING_LOGD("get SEI, ignore\n");
        return true;
    }

    if (!rtmpPacket_) {
        rtmpPacket_ = RtmpPacket::Create();
        // flags/not config/cts
        rtmpPacket_->SetCapacity(5);
    }

    uint8_t flag = FLV_CODEC_H264 | ((naluType == NAL_IDR ? FLV_KEY_FRAME : FLV_INTER_FRAME) << 4);
    rtmpPacket_->Assign(flag);
    rtmpPacket_->Append(0x01); // AVCPacketType, 0 sequence header, 1 AVC NALU

    uint8_t cts[3] = {0x00, 0x00, 0x00};
    rtmpPacket_->Append(cts, 3);

    uint32_t length = htonl((int)len);
    rtmpPacket_->Append((char *)&length, 4);

    rtmpPacket_->Append(data, (int)len);

    static std::chrono::system_clock::time_point start = std::chrono::system_clock::now();
    static uint64_t elapse = 0;
    elapse = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now() - start).count();

    rtmpPacket_->timestamp = (uint32_t)elapse; // dts
    rtmpPacket_->chunkId = CHUNK_VIDEO;
    rtmpPacket_->streamId = STREAM_MEDIA;
    rtmpPacket_->typeId = RTMPMessageType::VIDEO;
    rtmpPacket_->bodySize = rtmpPacket_->Size();
    if (callback_) {
        callback_(rtmpPacket_);
    }

    return false;
}

bool RtmpEncoderH264::InputFrame(const DataBuffer::Ptr &buffer)
{
    if (buffer && buffer->Size() > 0) {
        return InputFrame(buffer->Data(), buffer->Size());
    } else {
        return false;
    }
}

void RtmpEncoderH264::MakeConfigPacket()
{
    if (!sps_.empty() && !pps_.empty()) {
        MakeVideoConfigPkt();
    }
}

void RtmpEncoderH264::MakeVideoConfigPkt()
{
    SHARING_LOGD("trace.");
    if (sps_.empty()) {
        return;
    }

    uint8_t flag = FLV_CODEC_H264;
    flag |= (FLV_KEY_FRAME << 4);
    uint8_t cts[3] = {0x00, 0x00, 0x00};
    uint8_t version = 0x01;

    auto rtmpPkt = RtmpPacket::Create();
    rtmpPkt->SetCapacity(1024);
    // header
    rtmpPkt->Append(&flag, 1);
    rtmpPkt->Append(0x00); // AVCPacketType, 0 sequence header, 1 AVC NALU
    // cts
    rtmpPkt->Append(cts, 3);
    rtmpPkt->Append(version);
    rtmpPkt->Append(sps_.data(), 3); // profile-compat-level
    rtmpPkt->Append(0xff);           // 6 bits reserved + 2 bits nal size length - 1 (11)
    rtmpPkt->Append(0xe1);           // 3 bits reserved + 5 bits number of sps (00001)
    // sps
    uint16_t size = (uint16_t)sps_.size();
    size = htons(size);
    rtmpPkt->Append((uint8_t *)&size, 2);
    rtmpPkt->Append(sps_.data(), sps_.size());
    // pps
    rtmpPkt->Append(version);
    size = (uint16_t)pps_.size();
    size = htons(size);
    rtmpPkt->Append((uint8_t *)&size, 2);
    rtmpPkt->Append(pps_.data(), pps_.size());

    rtmpPkt->bodySize = rtmpPkt->Size();
    rtmpPkt->chunkId = CHUNK_VIDEO;
    rtmpPkt->streamId = STREAM_MEDIA;
    rtmpPkt->timestamp = 0;
    rtmpPkt->typeId = RTMPMessageType::VIDEO;

    if (callback_) {
        callback_(rtmpPkt);
    }
}
} // namespace Sharing
} // namespace OHOS